#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
PyTorch GPU验证脚本
安装PyTorch CUDA版本后运行此脚本验证是否成功
"""

import sys
import os

# 添加项目根目录到路径
project_root = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, project_root)

def verify_pytorch_gpu():
    """验证PyTorch GPU支持"""
    print("🎯 PyTorch GPU支持验证")
    print("=" * 60)
    
    try:
        import torch
        print(f"✅ PyTorch版本: {torch.__version__}")
        
        # 检查CUDA可用性
        cuda_available = torch.cuda.is_available()
        print(f"✅ CUDA可用性: {cuda_available}")
        
        if cuda_available:
            print("\n📱 CUDA详细信息:")
            print(f"  CUDA版本: {torch.version.cuda}")
            print(f"  GPU数量: {torch.cuda.device_count()}")
            print(f"  当前GPU: {torch.cuda.get_device_name()}")
            
            # 显存信息
            gpu_props = torch.cuda.get_device_properties(0)
            print(f"  GPU显存: {gpu_props.total_memory / (1024**3):.2f} GB")
            print(f"  计算能力: {gpu_props.major}.{gpu_props.minor}")
            
            # 测试GPU张量操作
            print("\n🧪 GPU张量操作测试:")
            try:
                # 创建张量并移动到GPU
                x = torch.randn(1000, 1000)
                print(f"  CPU张量创建成功: {x.shape}")
                
                x_gpu = x.cuda()
                print(f"  GPU张量移动成功: {x_gpu.device}")
                
                # GPU计算
                y_gpu = torch.matmul(x_gpu, x_gpu)
                print(f"  GPU矩阵乘法成功: {y_gpu.shape}")
                
                # 检查结果是否在GPU上
                if y_gpu.is_cuda:
                    print("  ✅ GPU计算验证通过")
                else:
                    print("  ❌ GPU计算结果未在GPU上")
                    
                print("\n🎉 GPU支持验证完全成功!")
                return True
                
            except Exception as e:
                print(f"  ❌ GPU张量操作测试失败: {e}")
                return False
        else:
            print("\n❌ CUDA不可用")
            print("💡 请按照安装指南重新安装支持CUDA的PyTorch版本")
            return False
            
    except ImportError as e:
        print(f" 无法导入PyTorch: {e}")
        print("💡 请先安装PyTorch")
        return False
    except Exception as e:
        print(f" 验证过程中出错: {e}")
        return False

def test_plw_model_gpu():
    """测试排列5模型GPU支持"""
    print("\n🎯 排列5模型GPU支持测试")
    print("=" * 60)
    
    try:
        import torch
        from algorithms.plw_sequence_lstm import PLWSequenceLSTM
        
        # 创建模型
        print(" 创建PLW序列LSTM模型...")
        model = PLWSequenceLSTM(input_dim=5, hidden_dim=64, num_layers=2, dropout=0.2)
        print(f"  ✅ 模型创建成功")
        print(f"  参数数量: {sum(p.numel() for p in model.parameters()):,}")
        
        if torch.cuda.is_available():
            # 将模型移动到GPU
            print(" 将模型移动到GPU...")
            model = model.cuda()
            print(f"  ✅ 模型已移动到: {next(model.parameters()).device}")
            
            # 创建测试数据
            print("📥 创建测试数据...")
            batch_size, seq_len, input_dim = 4, 5, 5
            test_input = torch.randint(0, 10, (batch_size, seq_len, input_dim)).float().cuda()
            print(f"  ✅ 测试数据创建成功: {test_input.shape} on {test_input.device}")
            
            # 前向传播测试
            print("⚡ 执行前向传播...")
            model.eval()
            with torch.no_grad():
                output = model(test_input)
                print(f"  ✅ 前向传播成功: {output.shape} on {output.device}")
            
            print("🎉 PLW模型GPU支持测试通过!")
            return True
        else:
            print("⚠️ CUDA不可用，跳过GPU测试")
            return False
            
    except Exception as e:
        print(f" PLW模型GPU支持测试失败: {e}")
        import traceback
        print(traceback.format_exc())
        return False

def show_gpu_training_benefits():
    """显示GPU训练的优势"""
    print("\n🚀 GPU训练优势")
    print("=" * 60)
    
    benefits = [
        "⚡ 训练速度提升5-10倍",
        " 支持更大的batch_size",
        "📈 可以训练更复杂的模型",
        "⏱️ 减少训练时间从小时级到分钟级",
        "🔋 更好的内存管理",
        "🔄 支持混合精度训练",
        "🎯 多GPU并行训练支持"
    ]
    
    for benefit in benefits:
        print(f"  {benefit}")

if __name__ == "__main__":
    print("🚀 开始PyTorch GPU支持验证")
    print("🎯 验证安装的PyTorch是否支持CUDA")
    print()
    
    # 验证PyTorch GPU支持
    gpu_support = verify_pytorch_gpu()
    
    # 测试PLW模型GPU支持
    plw_gpu_support = test_plw_model_gpu()
    
    # 显示GPU训练优势
    show_gpu_training_benefits()
    
    print("\n" + "=" * 80)
    print("验证结果:")
    print(f"  PyTorch GPU支持: {'✅ 通过' if gpu_support else '❌ 失败'}")
    print(f"  PLW模型GPU支持: {'✅ 通过' if plw_gpu_support else '❌ 失败'}")
    
    if gpu_support and plw_gpu_support:
        print("\n🎉 恭喜！PyTorch GPU支持已成功配置！")
        print("🎯 排列5序列LSTM模型训练现在将自动使用GPU")
        print("🚀 训练速度将大幅提升！")
    else:
        print("\n⚠️ GPU支持配置存在问题")
        print("💡 请按照安装指南重新安装PyTorch")